#include <hip/hip_runtime.h>
#include <iostream>
#include <vector>
#include <cassert>
#include <cmath>

// 错误检查宏
#define HIP_CHECK(cmd) \
    do { \
        hipError_t error = (cmd); \
        if (error != hipSuccess) { \
            std::cerr << "HIP Error: " << hipGetErrorString(error) \
                      << " at " << __FILE__ << ":" << __LINE__ << std::endl; \
            exit(EXIT_FAILURE); \
        } \
    } while (0)

// 控制变量：是否释放资源
static bool g_cleanup_resources = true;

// 数据处理核函数
__global__ void dataProcessingKernel(float* input, float* output, int size, float scale) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        // 模拟真实数据处理：缩放并添加偏移
        output[idx] = input[idx] * scale + idx * 0.01f;
    }
    return;
}

// 主机回调函数 - 模拟真实场景中的后处理
void postProcessCallback(void* userData) {
    int* processedCount = static_cast<int*>(userData);
    (*processedCount)++;

    // 模拟真实后处理操作
    if (*processedCount % 10 == 0) {
        std::cout << "Processed " << *processedCount << " batches" << std::endl;
    }
    return;
}

// 用户对象析构回调
void userObjectDestructor(void* userObject) {
    std::cout << "User object destructor called" << std::endl;
    return;
}

// 创建并执行HIP Graph，包含用户对象管理测试
void runGraphProcessing(int dataSize) {
    hipGraph_t graph = nullptr;
    hipGraphExec_t graphExec = nullptr;
    hipStream_t stream = nullptr;

    float *d_input = nullptr, *d_output = nullptr;
    float *h_input = nullptr, *h_output = nullptr;

    // 处理计数（用于回调函数）
    int processedCount = 0;

    // 创建用户对象
    hipUserObject_t userObject = nullptr;
    int userData = 12345; // 模拟用户数据

    std::cout << "=== Testing HIP Graph with User Object Management ===" << std::endl;

    // 创建用户对象
    HIP_CHECK(hipUserObjectCreate(&userObject,
                                 &userData,
                                 userObjectDestructor,
                                 1, // 初始引用计数
                                 hipUserObjectNoDestructorSync));
    std::cout << "Created user object" << std::endl;

    // 创建流和图
    HIP_CHECK(hipStreamCreate(&stream));
    HIP_CHECK(hipGraphCreate(&graph, 0));

    std::cout << "Creating HIP Graph with data size: " << dataSize << std::endl;

    // 1. 设备内存分配节点 - 输入
    hipMemAllocNodeParams allocParams = {};
    allocParams.bytesize = dataSize * sizeof(float);
    allocParams.poolProps.allocType = hipMemAllocationTypePinned;
    allocParams.poolProps.location.type = hipMemLocationTypeDevice;
    allocParams.poolProps.location.id = 0;

    hipGraphNode_t allocNode = nullptr;
    HIP_CHECK(hipGraphAddMemAllocNode(&allocNode, graph, nullptr, 0, &allocParams));
    d_input = reinterpret_cast<float*>(allocParams.dptr);

    // 2. 设备内存分配节点 - 输出
    hipMemAllocNodeParams outputAllocParams = {};
    outputAllocParams.bytesize = dataSize * sizeof(float);
    outputAllocParams.poolProps.allocType = hipMemAllocationTypePinned;
    outputAllocParams.poolProps.location.type = hipMemLocationTypeDevice;
    outputAllocParams.poolProps.location.id = 0;

    hipGraphNode_t outputAllocNode = nullptr;
    HIP_CHECK(hipGraphAddMemAllocNode(&outputAllocNode, graph, nullptr, 0, &outputAllocParams));
    d_output = reinterpret_cast<float*>(outputAllocParams.dptr);

    // 3. 主机内存分配（图外）
    h_input = new float[dataSize];
    h_output = new float[dataSize];

    // 初始化输入数据
    for (int i = 0; i < dataSize; i++) {
        h_input[i] = static_cast<float>(i) * 0.1f;
    }

    // 4. 主机到设备拷贝节点
    hipGraphNode_t h2dNode = nullptr;
    hipMemcpy3DParms h2dParams = {};
    h2dParams.srcPtr = make_hipPitchedPtr(h_input, dataSize * sizeof(float), dataSize, 1);
    h2dParams.dstPtr = make_hipPitchedPtr(d_input, dataSize * sizeof(float), dataSize, 1);
    h2dParams.extent = make_hipExtent(dataSize * sizeof(float), 1, 1);
    h2dParams.kind = hipMemcpyHostToDevice;
    HIP_CHECK(hipGraphAddMemcpyNode(&h2dNode, graph, nullptr, 0, &h2dParams));

    // 5. 核函数节点
    const int blockSize = 256;
    const int gridSize = (dataSize + blockSize - 1) / blockSize;

    // 创建可修改的变量用于内核参数
    int kernelSize = dataSize;
    float kernelScale = 2.5f;

    hipGraphNode_t kernelNode = nullptr;
    hipKernelNodeParams kernelParams = {};

    // 正确构建内核参数数组
    void* kernelArgs[4];
    kernelArgs[0] = &d_input;
    kernelArgs[1] = &d_output;
    kernelArgs[2] = &kernelSize;
    kernelArgs[3] = &kernelScale;

    kernelParams.func = reinterpret_cast<void*>(dataProcessingKernel);
    kernelParams.gridDim = dim3(gridSize);
    kernelParams.blockDim = dim3(blockSize);
    kernelParams.sharedMemBytes = 0;
    kernelParams.extra = nullptr;
    kernelParams.kernelParams = kernelArgs;
    HIP_CHECK(hipGraphAddKernelNode(&kernelNode, graph, &h2dNode, 1, &kernelParams));

    // 6. 设备到主机拷贝节点
    hipGraphNode_t d2hNode = nullptr;
    hipMemcpy3DParms d2hParams = {};
    d2hParams.srcPtr = make_hipPitchedPtr(d_output, dataSize * sizeof(float), dataSize, 1);
    d2hParams.dstPtr = make_hipPitchedPtr(h_output, dataSize * sizeof(float), dataSize, 1);
    d2hParams.extent = make_hipExtent(dataSize * sizeof(float), 1, 1);
    d2hParams.kind = hipMemcpyDeviceToHost;
    HIP_CHECK(hipGraphAddMemcpyNode(&d2hNode, graph, &kernelNode, 1, &d2hParams));

    // 7. 主机回调节点
    hipGraphNode_t hostNode = nullptr;
    hipHostNodeParams hostParams = {};
    hostParams.fn = postProcessCallback;
    hostParams.userData = &processedCount;
    HIP_CHECK(hipGraphAddHostNode(&hostNode, graph, &d2hNode, 1, &hostParams));

    // 8. 保留用户对象到图中
    HIP_CHECK(hipGraphRetainUserObject(graph, userObject, 1, hipGraphUserObjectMove));
    std::cout << "Retained user object in graph" << std::endl;

    // 9. 设备内存释放节点（根据控制变量决定是否添加）
    hipGraphNode_t freeInputNode = nullptr, freeOutputNode = nullptr;
    if (g_cleanup_resources) {
        HIP_CHECK(hipGraphAddMemFreeNode(&freeInputNode, graph, &hostNode, 1, d_input));
        HIP_CHECK(hipGraphAddMemFreeNode(&freeOutputNode, graph, &freeInputNode, 1, d_output));
    }

    // 实例化并执行图
    std::cout << "Instantiating and executing graph..." << std::endl;
    HIP_CHECK(hipGraphInstantiate(&graphExec, graph, nullptr, nullptr, 0));
    HIP_CHECK(hipGraphLaunch(graphExec, stream));
    HIP_CHECK(hipStreamSynchronize(stream));

    // 验证结果
    bool success = true;
    for (int i = 0; i < dataSize; i++) {
        float expected = h_input[i] * kernelScale + i * 0.01f;
        if (std::fabs(h_output[i] - expected) > 0.001f) {
            std::cout << "Mismatch at index " << i << ": expected " << expected
                      << ", got " << h_output[i] << std::endl;
            success = false;
            break;
        }
    }
    std::cout << "Graph execution " << (success ? "SUCCESS" : "FAILED") << std::endl;

    // 资源清理
    if (g_cleanup_resources) {
        // 从图中释放用户对象
        HIP_CHECK(hipGraphReleaseUserObject(graph, userObject, 1));
        std::cout << "Released user object from graph" << std::endl;

        // 释放用户对象
        HIP_CHECK(hipUserObjectRelease(userObject, 1));
        std::cout << "Released user object" << std::endl;

        // 释放其他资源
        if (graphExec) {
            HIP_CHECK(hipGraphExecDestroy(graphExec));
        }
        if (graph) {
            HIP_CHECK(hipGraphDestroy(graph));
        }
        if (stream) {
            HIP_CHECK(hipStreamDestroy(stream));
        }
        std::cout << "All graph resources cleaned up" << std::endl;
    } else {
        std::cout << "Resources retained (leak test)" << std::endl;
        std::cout << "User objects retained (leak test)" << std::endl;
    }

    // 主机内存泄漏不是测试目标，默认释放
    delete[] h_input;
    delete[] h_output;
    std::cout << "=== Test Completed ===" << std::endl;
}

int main(int argc, char** argv) {
    // 解析命令行参数
    for (int i = 1; i < argc; i++) {
        if (std::string(argv[i]) == "--leak") {
            g_cleanup_resources = false;
            std::cout << "Leak test mode enabled" << std::endl;
        }
    }

    // 检查设备
    int deviceCount = 0;
    HIP_CHECK(hipGetDeviceCount(&deviceCount));
    if (deviceCount == 0) {
        std::cerr << "No HIP devices found" << std::endl;
        return 1;
    }
    std::cout << "Found " << deviceCount << " HIP device(s)" << std::endl;

    // 设置设备
    HIP_CHECK(hipSetDevice(0));

    // 添加设备内存检查
    size_t freeMemBefore = 0, totalMemBefore = 0;
    size_t freeMemAfter = 0, totalMemAfter = 0;

    HIP_CHECK(hipMemGetInfo(&freeMemBefore, &totalMemBefore));

    // 获取设备属性
    hipDeviceProp_t prop;
    HIP_CHECK(hipGetDeviceProperties(&prop, 0));
    std::cout << "Using device: " << prop.name << std::endl;

    // 运行图处理
    runGraphProcessing(4096);  // 处理4096个元素

    HIP_CHECK(hipMemGetInfo(&freeMemAfter, &totalMemAfter));

    if (g_cleanup_resources) {
        size_t leakedDeviceMemory = freeMemBefore - freeMemAfter;
        if (leakedDeviceMemory > 1024) { // 1KB 阈值
            std::cout << "WARNING: Possible device memory leak: "
                      << leakedDeviceMemory << " bytes" << std::endl;
        }
    }

    return 0;
}